import re

import chromadb
from chromadb.utils import embedding_functions
from docx import Document
import os
from openai import OpenAI

# 初始化客户端
# chroma_client = chromadb.Client()
chroma_client = chromadb.PersistentClient(path="./chroma_db")


# 1:读取word文件:pip install python-docx
def load_employee_file(file_path):
    """
    这段代码的功能是读取Word文档并提取其中的文本内容。具体逻辑如下：

1. 使用`Document`类加载指定路径的Word文件
2. 遍历文档中的所有段落，提取每段的文本内容并去除首尾空白字符
3. 将处理后的文本添加到列表中
4. 最后将所有文本用换行符连接成一个字符串并返回

注释部分提到的读取表格、图片和去除特殊符号功能尚未实现。

    """
    document = Document(file_path)
    all_text = []
    for para in document.paragraphs:
        clean_text = para.text.strip()
        # print(clean_text)
        # 读取表格,图片
        # 去除特殊符号
        # 去除特殊符号，只保留中文、英文、数字和基本标点符号
        clean_text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\.,;:!?()""''\-]', '', clean_text)
        # 去除多余的空白字符
        clean_text = re.sub(r'\s+', ' ', clean_text)
        all_text.append(clean_text)

    # print(all_text)

    return '\n'.join(all_text)


# 简易文本切分：按固定长度切分（避免片段太长/太短）
def split_text(text, chunk_size=200):
    """
    这段代码定义了一个文本分割函数，功能是将输入的文本按指定大小切分成多个块：

1. `chunks = []` - 创建空列表存储文本块
2. `for i in range(0, len(text), chunk_size):` - 以chunk_size为步长遍历文本
3. `chunks.append(text[i:i + chunk_size])` - 截取文本片段添加到列表
4. `return chunks` - 返回分割后的文本块列表

注意：函数名"spilt_text"拼写有误，应为"split_text"。

这种简单的文本分割方法存在以下几个弊端：
语义不完整：按固定字符数切割可能会将完整的句子或段落切断，破坏语义连贯性
关键信息分离：可能将标题与内容、问题与答案等关联信息分割到不同的块中
上下文缺失：检索时可能只获取到片段的一部分，缺乏足够的上下文理解
重叠信息处理：重要信息如果跨越两个块，可能导致检索时无法完整获取
语言结构忽视：没有考虑句子边界、段落结构等自然语言特征
更好的方法是使用语义感知的分割策略，如按句子边界分割、保持段落完整或使用专门的文本分割工具（如LangChain的RecursiveCharacterTextSplitter）。
    """
    chunks = []
    for i in range(0, len(text), chunk_size):
        # print(i)
        # print(i + chunk_size)
        # print(text[i:i + chunk_size])
        chunks.append(text[i:i + chunk_size])
    return chunks


def spilt_text_1(text, chunk_size=200):
    """
    基于语义的文本分割函数，按句子边界切分文本并组合成适当大小的块
    """
    # 使用语义分割（按句子分割）
    try:
        import spacy
        # 加载中文语言模型
        nlp = spacy.load("zh_core_web_sm")
        # 处理文本
        doc = nlp(text)
        # 按句子分割
        sentences = [sent.text.strip() for sent in doc.sents if sent.text.strip()]

        # 将句子组合成适当大小的块
        chunks = []
        current_chunk = []
        current_length = 0

        for sentence in sentences:
            sentence_length = len(sentence)
            # 如果当前块加上新句子超过chunk_size，或者句子本身很长
            if current_length + sentence_length > chunk_size and current_chunk:
                # 保存当前块
                chunks.append(' '.join(current_chunk))
                # 开始新块
                current_chunk = [sentence]
                current_length = sentence_length
            else:
                # 添加句子到当前块
                current_chunk.append(sentence)
                current_length += sentence_length

        # 添加最后一个块
        if current_chunk:
            chunks.append(' '.join(current_chunk))

        return chunks
    except Exception as e:
        # 如果spaCy不可用或出错，回退到原有的固定长度分割
        print(f"使用固定长度分割，因为语义分割出错: {e}")
        chunks = []
        for i in range(0, len(text), chunk_size):
            chunks.append(text[i:i + chunk_size])
        return chunks


def spilt_text_2(text, chunk_size=200):
    """
    基于段落的语义分割函数
    """
    # 按段落分割（两个换行符）
    paragraphs = text.split('\n\n')

    chunks = []
    current_chunk = []
    current_length = 0

    for paragraph in paragraphs:
        paragraph = paragraph.strip()
        if not paragraph:
            continue

        paragraph_length = len(paragraph)
        # 如果当前块加上新段落超过chunk_size
        if current_length + paragraph_length > chunk_size and current_chunk:
            chunks.append('\n\n'.join(current_chunk))
            current_chunk = [paragraph]
            current_length = paragraph_length
        else:
            current_chunk.append(paragraph)
            current_length += paragraph_length

    # 添加最后一个块
    if current_chunk:
        chunks.append('\n\n'.join(current_chunk))

    # 如果仍然有块太大，再进行句子级别的分割
    final_chunks = []
    for chunk in chunks:
        if len(chunk) <= chunk_size:
            final_chunks.append(chunk)
        else:
            # 对大块进行句子级别的分割
            final_chunks.extend(split_by_sentences(chunk, chunk_size))

    return final_chunks


def split_by_sentences(text, chunk_size):
    """辅助函数：按句子分割文本"""
    try:
        import re
        # 简单的中文句子分割（按句号、问号、感叹号分割）
        sentences = re.split(r'[。！？]', text)
        sentences = [s.strip() for s in sentences if s.strip()]

        chunks = []
        current_chunk = []
        current_length = 0

        for sentence in sentences:
            if current_length + len(sentence) > chunk_size and current_chunk:
                chunks.append(''.join(current_chunk))
                current_chunk = [sentence]
                current_length = len(sentence)
            else:
                current_chunk.append(sentence)
                current_length += len(sentence)

        if current_chunk:
            chunks.append(''.join(current_chunk))
        return chunks
    except:
        # 简单的字符分割作为备选
        chunks = []
        for i in range(0, len(text), chunk_size):
            chunks.append(text[i:i + chunk_size])
        return chunks


def text_embeddings(texts):
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 如果您没有配置环境变量，请在此处用您的API Key进行替换
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"  # 百炼服务的base_url
    )

    # 如果输入是字符串，转换为列表
    if isinstance(texts, str):
        texts = [texts]

    # 阿里云百炼服务的text-embedding-v4模型批处理大小不能超过10
    batch_size = 10
    embeddings = []

    # 分批处理输入文本
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,  # 指定向量维度（仅 text-embedding-v3及 text-embedding-v4支持该参数）
            encoding_format="float"
        )

        # 提取嵌入向量
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chroma(texts):
    # 计算嵌入向量
    embeddings = text_embeddings(texts)

    # 创建嵌入函数（用于查询时保持一致性）
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 创建或获取集合（collection）
    collection = chroma_client.create_collection(
        name="my_collection",
        embedding_function=embedding_function
    )

    # 添加文档和预计算的嵌入向量到集合
    ids = [f"text_{i}" for i in range(len(texts))]
    collection.add(
        documents=texts,
        embeddings=embeddings,  # 使用预计算的嵌入向量
        ids=ids
    )


def show_chroma_collection():
    # 获取集合
    collection = chroma_client.get_collection(name="my_collection")

    # 获取集合中的所有数据
    data = collection.get()
    print(data)

    # 显示文档和ID
    for i in range(len(data['ids'])):
        print(f"ID: {data['ids'][i]}")
        print(f"文档: {data['documents'][i]}")

    # 获取向量数量
    count = collection.count()
    print(f"总向量数: {count}")


def query_chroma(query_text):
    # 创建嵌入函数
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 获取集合时指定嵌入函数
    collection = chroma_client.get_collection(
        name="my_collection",
        embedding_function=embedding_function
    )

    # 查询与特定文本相似的向量
    results = collection.query(
        query_texts=[query_text],
        n_results=5  # 返回最相似的5个结果
    )
    # {'ids': [['text_41', 'text_33', 'text_37', 'text_35', 'text_44']], 'embeddings': None, 'documents': [['年假（5-15天，按工作年限）、病假（按国家规定）、婚假/产假/陪产假等法定假期\n健康福利：每年1次', '到岗\n相应调整下班时间，保证8小时工作时长\n需要提前向直属主管申请并获得批准\n特殊工作安排\n项目紧急', '分钟以上\n早退：提前离开工作岗位\n旷工：未经批准不到岗工作\n考勤处理\n迟到3次以内：口头提醒\n迟到3', '要求\n员工应按时上下班，严格遵守考勤制度\n使用公司考勤系统进行打卡签到\n因故不能按时到岗应提前请假\n', '、年度旅游/聚餐、员工家庭日、兴趣社团支持\n第六章 培训发展\n6.1 培训体系\n新员工培训（3周）\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[None, None, None, None, None]], 'distances': [[0.6660666465759277, 0.8344534039497375, 0.9031039476394653, 0.9283781051635742, 0.9392238855361938]]}

    print(results)


def rag_answer_question(question, top_k=3):
    """
    RAG问答核心函数：
    1. 问题向量化 2. 检索相似片段 3. 结合片段通过LLM生成答案
    """

    # 创建嵌入函数
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 获取集合时指定嵌入函数
    collection = chroma_client.get_collection(
        name="my_collection",
        embedding_function=embedding_function
    )

    # 查询与特定文本相似的向量
    results = collection.query(
        query_texts=[question],
        n_results=5  # 返回最相似的5个结果
    )
    # [['年假（5-15天，按工作年限）、病假（按国家规定）、婚假/产假/陪产假等法定假期\n健康福利：每年1次', '到岗\n相应调整下班时间，保证8小时工作时长\n需要提前向直属主管申请并获得批准\n特殊工作安排\n项目紧急', '分钟以上\n早退：提前离开工作岗位\n旷工：未经批准不到岗工作\n考勤处理\n迟到3次以内：口头提醒\n迟到3', '要求\n员工应按时上下班，严格遵守考勤制度\n使用公司考勤系统进行打卡签到\n因故不能按时到岗应提前请假\n', '、年度旅游/聚餐、员工家庭日、兴趣社团支持\n第六章 培训发展\n6.1 培训体系\n新员工培训（3周）\n']]
    retrieved_chunks = list(set(results["documents"][0]))
    print("retrieved_chunks", retrieved_chunks)

    # 构建上下文（把片段拼起来，告诉LLM"基于这些内容回答"）
    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    print(f"\n【检索到的参考内容】\n{context}\n")
    prompt = f"""
    你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题。

    【员工手册内容】：
    {context}

    【用户问题】：
    {question}

    【回答要求】：
    1. 严格基于员工手册内容回答，不要编造或推测
    2. 如果手册中没有相关信息，请明确说明"手册中未提及此信息"
    3. 回答应简洁明了，条理清晰
    4. 如有多个要点，请分点列出
    5. 如果涉及具体条款，可引用相关内容
    """

    print("提示词:", prompt)

    messages = [{"role": "user", "content": prompt}]
    ai_response = get_response(messages)
    print("大模型返回的答案:", ai_response)


def get_response(messages):
    client = OpenAI(
        # 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：api_key="sk-xxx",
        api_key="sk-",
        # base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        base_url="http://localhost:11434/v1",

    )

    completion = client.chat.completions.create(
        model="qwen:latest",
        # model="qwen-plus",
        messages=messages,
    )
    return completion.choices[0].message.content


if __name__ == '__main__':
    result_text = load_employee_file('employee_manual.docx')
    # # print(result_text)
    chunks = spilt_text_1(result_text, chunk_size=50)
    print(chunks)
    # store_in_chroma(chunks)
    # print(chunks)
    # embeddings = text_embeddings(chunks)
    # print(embeddings)
    #
    # store_in_chroma(chunks,embeddings)
    # embeddings = text_embeddings("迟到3次以内：口头提醒")
    # print(embeddings)

    # show_chroma_collection()
    # query_chroma("员工工作满3年，年假有几天？需要提前多久申请")
    rag_answer_question("员工工作满3年，年假有几天？需要提前多久申请")
